#include "interpret_output.h"


inline float max_float(float a, float b) {
	return a>b ? a : b;
}

inline float min_float(float a, float b) {
	return a<b ? a : b;
}


inline float iou(box box1, box box2) {
	if (box1.class_index != box2.class_index) return 0;
        float tb = min_float(box1.x + 0.5*box1.w, box2.x + 0.5*box2.w) - max_float(box1.x - 0.5*box1.w, box2.x - 0.5*box2.w);
        float lr = min_float(box1.y + 0.5*box1.h, box2.y + 0.5*box2.h) - max_float(box1.y - 0.5*box1.h, box2.y - 0.5*box2.h);
        float intersection = 0;
	if (lr > 0 && tb > 0)
		intersection = lr * tb;
	return intersection / (box1.w * box1.h + box2.w * box2.h - intersection);
}

int argsort(int n, const box* x, int* indices)
{
 int i, j;
 int z = 0;
 for (i = 0; i<n; i++){
	 if(x[i].p!=0){
		 indices[z] = i;
		 z++;
	 }
 }
 for (i = 0; i < z; i++) {
  for (j = i + 1; j < z; j++)
  {
   if (x[indices[j]].p > x[indices[i]].p)
   {
    int index_tmp = indices[i];
    indices[i] = indices[j];
    indices[j] = index_tmp;
   }
  }
 }
 return z;
}

inline float get_value_from_4D_tensor(float* x, int row, int col, int box_loop,
	int index, int cell_H, int cell_W, int boxes_per_cell, int num) {
	int c = box_loop * num + index;
	//return x[col + cell_W * (row + cell_H * c)];
	return x[c + 128*(col + cell_W * row)];
}

inline float sigmoid(float x) {
	return 1.0f / (1 + expf(-x));
}

inline int softmax(float* x, int n) {
	int i, max_i;
	float max_x, sum;
	max_x = x[0];
	max_i = 0;
	for (i = 1; i < n; i++) {
		if (max_x < x[i]) {
			max_x = x[i];
			max_i = i;
		}
	}
	sum = 0;
	for (i = 0; i < n; i++) {
		x[i] = expf(x[i] - max_x);
		sum += x[i];
	}
	for (i = 0; i < n; i++)
		x[i] /= sum;
	return max_i;
}

Rect* interpret_output(float* output,
	int cell_H, int cell_W,
	int image_height, int image_width, int original_height, int original_width, int h_step, int w_step,
	int num_class, int boxes_per_cell,
        float threshold, float iou_threshold,
	int* numBoxesOut) {
	int i, j;
	int channels = num_class + 5;
        //float anchors[10] = { 1.08f, 1.19f, 3.42f, 4.41f, 6.63f, 11.38f, 9.42f, 5.11f, 16.62f, 10.52f };
        float anchors[10] = {7.00f, 8.75f,  13.56f,  18.08f,  14.43f,  71.75f,  22.31f, 111.42f,  35.88f, 173.83f};
	int total_box_num = cell_H * cell_W * boxes_per_cell;
	//int total_probs_num = total_box_num * num_class;
	box* boxes = (box*)calloc(total_box_num, sizeof(box));  //free
	float* probs = (float*)calloc(num_class, sizeof(float));  //free

	int row, col, box_loop, class_loop, box_index, class_index;
	float center_x, center_y, box_w, box_h, p_has_object, p_class;
	for (row = 0; row < cell_H; row++) {
		for (col = 0; col < cell_W; col++) {
			for (box_loop = 0; box_loop < boxes_per_cell; box_loop++) {
				box_index = box_loop + boxes_per_cell * (col + row * cell_W);

				p_has_object = get_value_from_4D_tensor(output, row, col, box_loop, 4,
					cell_H, cell_W, boxes_per_cell, channels);
				p_has_object = sigmoid(p_has_object);
				if (p_has_object >= threshold) {
					for (class_loop = 0; class_loop < num_class; class_loop++) {
						p_class = get_value_from_4D_tensor(output, row, col, box_loop, 5 + class_loop,
							cell_H, cell_W, boxes_per_cell, channels);
						probs[class_loop] = p_class;
					}
					class_index = softmax(probs, num_class);
					//printf("%d  %d  %d  %d  %f\n", row, col, box_loop, class_index, probs[class_index]);
					p_class = p_has_object * probs[class_index];
					boxes[box_index].class_index = class_index;
					if (p_class >= threshold){
                                            boxes[box_index].p = p_class;
                                            center_x = get_value_from_4D_tensor(output, row, col, box_loop, 0,
                                                    cell_H, cell_W, boxes_per_cell, channels);
                                            center_x = (col + sigmoid(center_x)) * image_width / cell_W;
                                            boxes[box_index].x = center_x;

                                            center_y = get_value_from_4D_tensor(output, row, col, box_loop, 1,
                                                    cell_H, cell_W, boxes_per_cell, channels);
                                            center_y = (row + sigmoid(center_y)) * image_height / cell_H;
                                            boxes[box_index].y = center_y;

                                            box_w = get_value_from_4D_tensor(output, row, col, box_loop, 2,
                                                    cell_H, cell_W, boxes_per_cell, channels);
                                            box_w = expf(box_w) * anchors[2 * box_loop] * image_width / cell_W;
                                            boxes[box_index].w = box_w;

                                            box_h = get_value_from_4D_tensor(output, row, col, box_loop, 3,
                                                    cell_H, cell_W, boxes_per_cell, channels);
                                            box_h = expf(box_h) * anchors[2 * box_loop + 1] * image_height / cell_H;
                                            boxes[box_index].h = box_h;
                                        }
                                        else boxes[box_index].p = 0.0;
                                }
				else boxes[box_index].p = 0.0;
			}
		}
	}


	 int* indices = (int*)calloc(total_box_num, sizeof(int));  //free
	 int valid_num = argsort(total_box_num, boxes, indices);
	 for (i = 0; i<valid_num; i++) {
	  if (boxes[indices[i]].p == -1) continue;
	  for (j = i + 1; j<valid_num; j++) {
	   if (boxes[indices[j]].p == -1) continue;
	   if (iou(boxes[indices[i]], boxes[indices[j]]) > iou_threshold)
	    boxes[indices[j]].p = -1;
	  }
	 }

	free(indices);

	*numBoxesOut = 0;
	for (box_index = 0; box_index < total_box_num; box_index++) {
		if (boxes[box_index].p > threshold)
			(*numBoxesOut)++;
	}
	i = 0;
	float w_half, h_half;
	float x1, x2, y1, y2;
	Rect* results = (Rect*)calloc(*numBoxesOut, sizeof(Rect));  //free
	float x_scale = original_width * 1.0f / image_width;
	float y_scale = original_height * 1.0f / image_height;
	/*
	if (w_step > 0)
		x_scale = y_scale;
	else
		y_scale = x_scale;
	*/
	for (box_index = 0; box_index < total_box_num; box_index++) {
		if (boxes[box_index].p > threshold) {
			w_half = boxes[box_index].w / 2;
			h_half = boxes[box_index].h / 2;
			x1 = boxes[box_index].x - w_half;
			if (x1 < 0) x1 = 0;
			x2 = boxes[box_index].x + w_half;
			if (x2 >= image_width) x2 = (float)image_width - 1;
			y1 = boxes[box_index].y - h_half;
			if (y1 < 0) y1 = 0;
			y2 = boxes[box_index].y + h_half;
			if (y2 >= image_height) y2 = (float)image_height - 1;
			results[i].x1 = (x1 - w_step) * x_scale;
			results[i].x2 = (x2 - w_step) * x_scale;
			results[i].y1 = (y1 - h_step) * y_scale;
			results[i].y2 = (y2 - h_step) * y_scale;
			results[i].class_index = boxes[box_index].class_index;
			results[i].p = boxes[box_index].p;
			i++;
		}
	}
	free(boxes);
	free(probs);
	return results;
}
